import torch
import torch.nn as nn
import torchvision.transforms.functional as tf

class DoubleConv(nn.Module):
    def __init__(self, inChannel, outChannel):
        super(DoubleConv, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(inChannel, outChannel, kernel_size=(3, 3), stride=(1, 1), padding=1, bias=False),
            nn.BatchNorm2d(outChannel),
            nn.ReLU(inplace=True),
            nn.Conv2d(outChannel, outChannel, kernel_size=(3, 3), stride=(1, 1), padding=1, bias=False),
            nn.BatchNorm2d(outChannel),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        return self.conv(x)

class Unet(nn.Module):
    def __init__(self, inChannel=3, outChannel=1, features=(64,128,256,512)):
        super(Unet, self).__init__()
        self.ups = nn.ModuleList()
        self.downs = nn.ModuleList()
        self.pool = nn.MaxPool2d(kernel_size=(2,2), stride=2)

        for feature in features:
            self.downs.append(DoubleConv(inChannel, feature))
            inChannel = feature

        for feature in reversed(features):
            self.ups.append(nn.ConvTranspose2d(feature*2, feature, kernel_size=(3, 3), stride=(1, 1), padding=(1,1)))
            self.ups.append(DoubleConv(inChannel=feature*2, outChannel=feature))

        self.bottleneck = DoubleConv(features[-1], features[-1]*2)

        self.finalConv = nn.Conv2d(features[0], outChannel, kernel_size=(3, 3), stride=(1, 1), padding=1, bias=False)

    def forward(self, x):
        copyAndCrops = []

        for down in self.downs:
            x=down(x)
            copyAndCrops.append(x)
            x=self.pool(x)

        x=self.bottleneck(x)
        copyAndCrops=copyAndCrops[::-1]

        for idx in range(0, len(self.ups), 2):
            x=self.ups[idx](x)
            copyAndCrop=copyAndCrops[idx//2]
            if copyAndCrop.shape != x.shape:
                x=tf.resize(x, size=copyAndCrop.shape[2:])

            concat = torch.cat((x,copyAndCrop), dim=1)
            x=self.ups[idx+1](concat)

        return self.finalConv(x)

def test():
    x=torch.randn((3,3,160,160))
    unet = Unet(inChannel=3, outChannel=1)
    preds = unet(x)
    print(x.shape)
    print(preds.shape)

if __name__ == "__main__":
    test()